import os
import subprocess
import shutil
import sys


def main(source_cwe_path, target_cwe_path):

    # 指定源CWE121文件夹路径和目标路径
    # source_cwe121_path = '/home/kali/Desktop/angr/C/testcases/CWE121_Stack_Based_Buffer_Overflow'
    # target_cwe121_path = '/home/kali/Desktop/angr/dataset/CWE121_Stack_Based_Buffer_Overflow_O3'  

    # your_source_cwe121_path ='/home/kokomi/static_project/test/C/testcases/CWE121_Stack_Based_Buffer_Overflow'                               
    # your_target_cwe121_path ='/home/kokomi/static_project/test/dataset/CWE121'                               # 填入你的目标路径

    # if your_source_cwe121_path :
    #     source_cwe121_path = your_source_cwe121_path

    # if your_target_cwe121_path :
    #     target_cwe121_path = your_target_cwe121_path

    source_cwe121_path = source_cwe_path
    target_cwe121_path = target_cwe_path


    bad_folder_suffix = 'bad'
    good_folder_suffix = 'good'

    # 遍历每个子文件夹
    for subdir in os.listdir(source_cwe121_path):
        source_dir_path = os.path.join(source_cwe121_path, subdir)
        target_dir_path = os.path.join(target_cwe121_path, subdir)

        # 确保是文件夹并且是需要处理的子文件夹
        if os.path.isdir(source_dir_path) and subdir.startswith('s'):
            makefile_path = os.path.join(source_dir_path, 'Makefile')

            # 检查Makefile是否存在
            if os.path.exists(makefile_path):
                # 修改Makefile以添加 -DOMITGOOD
                with open(makefile_path, 'r') as file:
                    lines = file.readlines()

                with open(makefile_path, 'w') as file:
                    for line in lines:
                        if '$(INCLUDE_MAIN)' in line and '$(CC)' in line:
                            line = line.replace(' -DOMITBAD -O3 -s -o $@', ' -DOMITGOOD -O3 -o $@')
                            
                        elif '$(INCLUDE_MAIN)' in line and '$(CPP)' in line:
                            line = line.replace(' -DOMITBAD -O3 -s -o $@', ' -DOMITGOOD -O3 -o $@')
                            
                        file.write(line)

                # 进入子文件夹并执行make命令
                subprocess.run(['make', 'clean'], cwd=source_dir_path)
                subprocess.run(['make', 'individuals'], cwd=source_dir_path)

                # 复制生成的.out文件到bad文件夹
                bad_folder_path = os.path.join(target_dir_path, bad_folder_suffix)
                os.makedirs(bad_folder_path, exist_ok=True)

                for file in os.listdir(source_dir_path):
                    if file.endswith('.out'):
                        shutil.copy(os.path.join(source_dir_path, file), bad_folder_path)

                # 修改Makefile以将 -DOMITGOOD 改为 -DOMITBAD
                with open(makefile_path, 'r') as file:
                    lines = file.readlines()

                with open(makefile_path, 'w') as file:
                    for line in lines:
                        line = line.replace(' -DOMITGOOD', ' -DOMITBAD')
                        file.write(line)
            
                # 进入子文件夹并执行make命令
                subprocess.run(['make', 'clean'], cwd=source_dir_path)
                subprocess.run(['make', 'individuals'], cwd=source_dir_path)

                # 创建good文件夹并复制.out文件到good文件夹
                good_folder_path = os.path.join(target_dir_path, good_folder_suffix)
                os.makedirs(good_folder_path, exist_ok=True)

                for file in os.listdir(source_dir_path):
                    if file.endswith('.out'):
                        shutil.copy(os.path.join(source_dir_path, file), good_folder_path)

    print("所有操作完成。")

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print("用法: python juliet.py <source_cwe_path> <target_cwe_path>")

    else:
        source_cwe_path = sys.argv[1]
        target_cwe_path = sys.argv[2]
        main(source_cwe_path, target_cwe_path)